#!/bin/bash

# 设置环境变量以使用合适的GPU
export CUDA_VISIBLE_DEVICES=0,1,2,3
export TOKENIZERS_PARALLELISM=false

# 项目根目录
ROOT_DIR=$(dirname $(dirname "$0"))
echo "项目根目录: $ROOT_DIR"

# 创建输出目录
OUTPUT_DIR="$ROOT_DIR/dataset/spatial_corr"
mkdir -p "$OUTPUT_DIR"
echo "输出目录: $OUTPUT_DIR"

# 步骤1: 处理空间相关性数据 - 将MAT文件转换为NPZ格式
echo "步骤1: 处理空间相关性数据..."
python "$ROOT_DIR/scripts/process_spatial_corr_data.py" \
    --input_dir="$ROOT_DIR/dy" \
    --output_dir="$OUTPUT_DIR" \
    --window_size=20 \
    --overlap=10

# 检查数据文件是否生成成功
if [ ! -f "$OUTPUT_DIR/spatial_correlation_data.npz" ]; then
    echo "错误: 数据文件未生成，请检查空间相关性数据处理脚本。"
    exit 1
fi

echo "数据处理完成，开始训练模型..."

# 步骤2: 使用Deepspeed启动多GPU训练
echo "步骤2: 使用Deepspeed训练模型..."

# 添加deepspeed参数
DEEPSPEED_ARGS="--num_nodes=1 \
               --num_gpus=4 \
               --master_addr=localhost \
               --master_port=29500"

# 启动训练
deepspeed $DEEPSPEED_ARGS "$ROOT_DIR/run_spatial_corr.py" \
    --task_name spatial_corr_prediction \
    --data SPATIAL_CORR \
    --root_path "$OUTPUT_DIR" \
    --data_path spatial_correlation_data.npz \
    --checkpoints "$ROOT_DIR/checkpoints" \
    --seq_len 20 \
    --label_len 10 \
    --pred_len 10 \
    --d_model 256 \
    --n_heads 8 \
    --d_ff 512 \
    --n_probes 32 \
    --llm_model Qwen3-4B \
    --llm_layers 6 \
    --batch_size 8 \
    --eval_batch_size 16 \
    --learning_rate 0.0001 \
    --train_epochs 100 \
    --patience 10 \
    --loss kl \
    --model_comment spatial_corr_only \
    --num_workers 4 \
    --use_transformer \
    --n_transformer_layers 4

echo "训练完成！" 